#!/usr/bin/python
# -*- coding: utf-8 -*-
from typing import List

import pandas as pd

from core.dataClasses import StockTradeDataColumnName
from core.enums import Period
from infrastructure.util import dateutils
from core.tactic import AbstractAppearTactic
from index.rsi import SmaRSI, SMA_RSI_COLUMN_PREFIX
from core.inventedtrader import TradeType

"""
思想：在股票市场上，买方和卖方力量的消长会影响股票的价格。如果股票的卖方力量大于卖方力量，则股票的价格会上涨；如果股票的卖方力量大于买方力量，则股票的价格会下跌。
RSI（Relative Strength Index）相对强弱指标，是用一种特定公式计算出来的值，投资者可以通过RSI的取值来判断股票的买入和卖出情况，进而预测未来股票的价格走势。
    RSI = 100 - 100/(1 + RS) 或 RSI = 100 * RS/(1 + RS)
    RS = UP/DOWN
    将RS = UP/DOWN带入RSI = 100 * RS/(1 + RS)可以推出：
        RSI = 100 * (UP/DOWN)/(1 + UP/DOWN) = 100 * UP/(UP + DOWN)
    其中：
        RSI表示相对强弱指标值；若t表示所参考数据的期数，UP表示t期内估价上涨幅度的平均值，DOWN表示t期内估价下跌幅度的平均值。
        在计算UP和DOWN的值时，需要求出上涨幅度的平均值和下跌幅度的平均值，一般使用算数平均值来求UP和DOWN。
            均值可分为简单移动平均值（SMA）、加权移动平均值（WMA）和指数移动平均值（EMA）。
策略：RSI的取值大于50越多，则表明股票上涨的力量超过下跌的力量更大。当RSI的取值小于50越多，则股票的下跌力量超过上涨力量更大。
RSI天数的差异：
    韦尔斯 维尔德（Wells Wilder）指出，通过运用月周期28日的一般来计算RSI的值进行预测是有效的，他推荐RSI的时间跨度默认值为14日。
    一些看盘软件设有6日RSI、12日RSI和24日RSI指标。6日近似一周的时间周期，12日可以看做半个月的时间跨度，而24日约为一个月的时间跨度。
超买线、超卖线、中心线
    RSI取值等于80或20分别为较常见的“超买线”和“超卖线”的刻画。RSI为80是股票超买的临界点，RSI为20是股票超卖的临界点，RSI取值为50设定为“中心线“，该线表明股票的买入力量等于卖出力量。
    当RSI大于80时，股票出现超买信号。股票买入力量过大，买入力量在未来可能会减小，所以股票未来价格可能会下跌，此时卖出股票，未来下跌后再买入股票，从而赚取价差。
    当RSI小于20时，股票出现超卖信号。股票卖出力量过大，卖出力量在未来可能会回归正常，所以股票未来价格可能会上涨，投资者此时可以做多股票，未来价格上涨后再卖出。
黄金交叉和死亡交叉：
    当短期RSI线向上穿过长期RSI线时，股票近期买入的力量较强，价格上涨的力量很大，其释放出一个较强的买入信号，这个信号被成为“黄金交叉”。
    当短期RSI线向下跌破长期RSI线时，股票近期卖出的力量较强，价格下跌的力量很大，其释放出一个较强的卖出信号，被成为“死亡交叉”。
"""


class SmaRSIOverdoseTactic(AbstractAppearTactic):
    """
    超卖、超买策略
    """

    def __init__(self, sell_line=80, buy_line=20, trade_types: List[TradeType] = [TradeType.BUY, TradeType.SELL], period = Period.day):
        super(SmaRSIOverdoseTactic, self).__init__()
        self.period = period
        self.trade_types = trade_types
        self.sell_line = sell_line
        self.buy_line = buy_line

    def setup(self, context):
        df: pd.DataFrame = context.data
        column = SMA_RSI_COLUMN_PREFIX + str(self.period)
        if column not in df.columns.values:
            SmaRSI(self.period, key= StockTradeDataColumnName.CLOSE).compute(context.data)
        if self.trade_types is None or len(self.trade_types)==0 or TradeType.BUY in self.trade_types:
            self.buy_date = df[df[column] <= self.buy_line].index.strftime(dateutils.DateFormat.Y_m_d).values
        if self.trade_types is None or len(self.trade_types)==0 or TradeType.SELL in self.trade_types:
            self.sell_date = df[df[column] <= self.sell_line].index.strftime(dateutils.DateFormat.Y_m_d).values


class SmaRSICrossTactic(AbstractAppearTactic):
    """
    金叉、死叉策略
    """
    def __init__(self, period1, period2, trade_types: List[TradeType] = [TradeType.BUY, TradeType.SELL], period = Period.day):
        super(SmaRSIOverdoseTactic, self).__init__()
        self.period = period
        self.trade_types = trade_types
        if period1 == period2:
            raise EOFError("period1 and period2 can`t be equal")
        if period1 < period2:
            self.period1 = period1
            self.period2 = period2
        else:
            self.period1 = period2
            self.period2 = period1

    def setup(self, context):
        df = context.data  # type:pd.DataFrame
        column1 = SMA_RSI_COLUMN_PREFIX + str(self.period1)
        column2 = SMA_RSI_COLUMN_PREFIX + str(self.period2)
        if column1 not in df.columns.values:
            SmaRSI(self.period1, key=StockTradeDataColumnName.CLOSE).compute(context.data)
        if column2 not in df.columns.values:
            SmaRSI(self.period2, key=StockTradeDataColumnName.CLOSE).compute(context.data)
        if self.order_type != TradeType.SELL:
            self.buy_date = df[
                (df[self.period1].shift(1) <= df[self.period2].shift(1))
                & (df[self.period1] > df[self.period2])
                ].index.strftime(dateutils.DateFormat.Y_m_d).values
        if self.order_type != TradeType.BUY:
            self.sell_date = df[
                (df[self.period1].shift(1) >= df[self.period2].shift(1))
                & (df[self.period1] < df[self.period2])
                ].index.strftime(dateutils.DateFormat.Y_m_d).values
